#!/usr/bin/env python3
"""
kNN.py
k-Nearest Neighbors
"""

import numpy as np

def kNN(d, data_lables, k=3, func_distance=lambda a, b: abs(a-b)):
    """
    Return a list containing the nereast k (data, lable) tuples.
    d: the data which is to determine its lable
    data_lables: a data:lable dict
    k: the nearest element number
    func_distance: function to compute the distance of data
    """
    neighbors = []
    neighbor_distances = []

    # sanity checks
    if len(data_lables) < k:
        raise Exception("kNN: not enough data numbers")

    # calculate the distances
    for data in data_lables.keys():
        neighbor_distances.append((data, func_distance(d, data)))
    neighbor_distances.sort(key=lambda item: item[1])

    neighbors = [(nb[0], data_lables[nb[0]]) for nb in neighbor_distances[:k]]
    return neighbors

def vote(neighbors):
    """
    Return the lable that most of the neighbors have.
    neighbors: list of (data, lable) tuples
    """
    lables = [nb[1] for nb in neighbors]
    lable_counts = {}
    for lable in lables:
        if not lable in lable_counts:
            lable_counts[lable] = 0
        else:
            lable_counts[lable] += 1
    L = list(lable_counts.items())
    L.sort(key=lambda item: item[1])
    return L[-1][0]

if __name__ == "__main__":
    test_data_lables = {1:'small', 2:'small', 10:'small',
                        20:'small', 30:'small', 40:'small',
                        60:'large', 70:'large', 80:'large',
                        90:'large', 100:'large'}
    d = 12
    k_neighbors = kNN(d, test_data_lables)
    print(k_neighbors)
    print("vote is '%s'" % vote(k_neighbors))

    test_data_lables = {(0, 0):'nobody', (1, 2):'nobody', (2, 2):'nobody',
                        (10, 10):'james', (20,20):'james', (30,30):'james'}
    def my_distance(a, b):
        return abs((a[0]-b[0])**2 + (a[1]-b[1])**2)
    d = (15, 20)
    k_neighbors = kNN(d, test_data_lables, func_distance=my_distance, k=2)
    print(k_neighbors)
    print("vote is '%s'" % vote(k_neighbors))
    
